#!/bin/env python
# -*- coding: utf-8 -*-
import taichi as ti
import numpy as np
from tqdm import tqdm
from scipy.spatial import cKDTree
from typing import List
from ..dither_util import DitherUtil
import cv2

@ti.data_oriented
class BayerDitherer(DitherUtil):
    """
    Bayer抖动算法实现类。
    支持灰度和彩色图像。
    """
    
    def __init__(self, arch=ti.cpu, debug=False, max_value=255):
        """
        初始化Bayer抖动处理器。
        
        :param arch: Taichi架构，默认为CPU，可以设置为ti.gpu启用GPU加速
        :param debug: 是否开启调试模式，默认为False
        :param max_value: 像素最大值，默认为255
        """
        super().__init__(arch, debug, max_value)
    
    def bit_reverse(self, x: int, n: int) -> int:
        """
        比特反转操作。
        
        :param x: 输入整数
        :param n: 位数
        :return: 反转后的值
        """
        return int(bin(x)[2:].zfill(n)[::-1], 2)
    
    def bit_interleave(self, x: int, y: int, n: int) -> int:
        """
        比特交织操作。
        
        :param x: 第一个输入整数
        :param y: 第二个输入整数
        :param n: 位数
        :return: 交织后的值
        """
        x = bin(x)[2:].zfill(n)
        y = bin(y)[2:].zfill(n)
        return int(''.join(''.join(i) for i in zip(x, y)), 2)
    
    def bayer_entry(self, x: int, y: int, n: int) -> int:
        """
        计算Bayer矩阵的条目值。
        
        :param x: x坐标
        :param y: y坐标
        :param n: 阶数
        :return: Bayer矩阵中的值
        """
        return self.bit_reverse(self.bit_interleave(x ^ y, y, n), 2*n)
    
    def bayer_matrix(self, n: int) -> List[List[int]]:
        """
        生成Bayer矩阵。
        
        :param n: 阶数
        :return: Bayer矩阵
        """
        r = range(2**n)
        return [[self.bayer_entry(x, y, n) for x in r] for y in r]
    
    @ti.kernel
    def dither_bayer(self, img: ti.types.ndarray(), bayerM: ti.types.ndarray()):
        """
        Bayer抖动算法实现。
        
        :param img: 输入图像数组
        :param bayerM: Bayer矩阵
        """
        h, w = img.shape
        n = bayerM.shape[0]
        for i, j in ti.ndrange(h, w):
            threshold = bayerM[i % n, j % n] * self.MAX // (n**2)
            
            # 固定两色处理 (黑白)
            if img[i, j] > threshold:
                img[i, j] = self.PALETTE[1]
            else:
                img[i, j] = self.PALETTE[0]
    
    # 性能优化的彩色Bayer抖动实现
    def apply_multi_bayer(self, img, bayerM, num_colors):
        """
        应用支持多色值的Bayer抖动算法。
        因为Taichi内核不支持动态列表索引，所以我们在非内核函数中处理。
        
        :param img: 输入图像numpy数组
        :param bayerM: Bayer矩阵numpy数组
        :param num_colors: 色值数量
        :return: 处理后的图像
        """
        h, w = img.shape
        n = bayerM.shape[0]
        
        # 转换为numpy数组处理
        img_np = img.copy()
        bayerM_np = bayerM.to_numpy() if hasattr(bayerM, 'to_numpy') else bayerM
        
        # 将调色板转换为numpy数组
        palette_np = np.array(self.PALETTE, dtype=np.int32)
        
        # 计算Bayer阈值因子，避免溢出
        bayer_factor = float(self.MAX) / float(n**2)
        
        for i in range(h):
            for j in range(w):
                # 计算阈值，避免整数溢出
                threshold = int(bayerM_np[i % n, j % n] * bayer_factor)
                
                # 获取当前像素
                pixel = img_np[i, j]
                
                # 找到最接近的调色板颜色
                if num_colors == 2:
                    # 黑白处理 - 简单阈值比较
                    if pixel > threshold:
                        img_np[i, j] = palette_np[1]
                    else:
                        img_np[i, j] = palette_np[0]
                else:
                    # 多色处理 - 根据调色板查找
                    # 缩放像素值到调色板范围
                    scaled_value = int((pixel + threshold) * num_colors / (self.MAX + 1))
                    idx = max(0, min(scaled_value, num_colors - 1))
                    img_np[i, j] = palette_np[idx]
                    
                    # 计算量化误差
                    quant_error = pixel - palette_np[idx]
                    
                    # 误差扩散
                    if j + 1 < w:
                        img_np[i, j + 1] = np.clip(img_np[i, j + 1] + quant_error * 7 // 16, 0, self.MAX)
                    if i + 1 < h:
                        if j - 1 >= 0:
                            img_np[i + 1, j - 1] = np.clip(img_np[i + 1, j - 1] + quant_error * 3 // 16, 0, self.MAX)
                        img_np[i + 1, j] = np.clip(img_np[i + 1, j] + quant_error * 5 // 16, 0, self.MAX)
                        if j + 1 < w:
                            img_np[i + 1, j + 1] = np.clip(img_np[i + 1, j + 1] + quant_error * 1 // 16, 0, self.MAX)
        
        return img_np
    
    # 性能优化的彩色Bayer抖动实现 - 使用Taichi内核加速
    @ti.kernel
    def apply_color_bayer_kernel(self, 
                              img_in: ti.types.ndarray(), 
                              img_out: ti.types.ndarray(),
                              bayerM: ti.types.ndarray(),
                              palette: ti.types.ndarray()):
        """
        使用Taichi内核实现的彩色Bayer抖动算法
        
        :param img_in: 输入彩色图像
        :param img_out: 输出彩色图像
        :param bayerM: Bayer矩阵
        :param palette: 颜色调色板，形状为(n, 3)
        """
        h, w = img_in.shape[0], img_in.shape[1]
        n = bayerM.shape[0]
        n_colors = palette.shape[0]
        
        for i, j in ti.ndrange(h, w):
            # 计算阈值，避免整数溢出
            threshold = int(bayerM[i % n, j % n] * self.MAX / (n**2))
            
            # 获取当前像素
            r = img_in[i, j, 0]
            g = img_in[i, j, 1]
            b = img_in[i, j, 2]
            
            # 找到最接近的调色板颜色
            min_dist = 1e10
            best_idx = 0
            
            for k in range(n_colors):
                pr = palette[k, 0]
                pg = palette[k, 1]
                pb = palette[k, 2]
                
                # 计算颜色距离
                dist = (r-pr)*(r-pr) + (g-pg)*(g-pg) + (b-pb)*(b-pb)
                
                if dist < min_dist:
                    min_dist = dist
                    best_idx = k
            
            # 输出最佳匹配的颜色
            img_out[i, j, 0] = palette[best_idx, 0]
            img_out[i, j, 1] = palette[best_idx, 1]
            img_out[i, j, 2] = palette[best_idx, 2]
    
    # 修改后的彩色Bayer抖动实现，使用上面的内核函数
    def apply_color_bayer(self, img, bayerM, color_palette):
        """
        应用彩色Bayer抖动算法。
        
        :param img: 输入彩色图像numpy数组 (h, w, 3)
        :param bayerM: Bayer矩阵numpy数组
        :param color_palette: 颜色调色板 [(r,g,b), ...]
        :return: 处理后的彩色图像
        """
        h, w, c = img.shape
        
        # 转换为numpy数组处理
        img_np = img.copy()
        bayerM_np = bayerM.to_numpy() if hasattr(bayerM, 'to_numpy') else bayerM
        
        # 将调色板转换为numpy数组
        palette_np = np.array(color_palette, dtype=np.int32)
        
        # 创建输出图像
        output_img = np.zeros_like(img_np)
        
        # 使用GPU加速的Taichi内核处理
        if self.arch == ti.gpu:
            # 转换为Taichi数组
            img_ti = ti.ndarray(shape=img_np.shape, dtype=ti.i32)
            img_ti.from_numpy(img_np.astype(np.int32))
            
            output_ti = ti.ndarray(shape=output_img.shape, dtype=ti.i32)
            
            palette_ti = ti.ndarray(shape=palette_np.shape, dtype=ti.i32)
            palette_ti.from_numpy(palette_np)
            
            # 执行GPU加速的内核函数
            self.apply_color_bayer_kernel(img_ti, output_ti, bayerM, palette_ti)
            
            # 确保最终图像只使用调色板中的颜色
            # 使用Taichi内核加速颜色映射
            self.map_to_color_palette(output_ti, palette_ti)
            
            # 转回numpy数组
            output_img = output_ti.to_numpy()
        else:
            # 回退到CPU处理方式
            # 计算Bayer阈值因子，避免溢出
            bayer_factor = self.MAX / (bayerM_np.shape[0]**2)
            
            # 使用更高效的颜色匹配算法
            palette_tree = cKDTree(palette_np)
            
            # 显示进度
            for i in tqdm(range(h), desc="应用Bayer抖动", unit="行"):
                for j in range(w):
                    # 计算阈值，避免整数溢出
                    threshold = int(bayerM_np[i % bayerM_np.shape[0], j % bayerM_np.shape[0]] * bayer_factor)
                    
                    # 获取当前像素
                    pixel = img_np[i, j].astype(np.int32)
                    
                    # 使用KD树查找最近颜色
                    _, idx = palette_tree.query(pixel)
                    best_color = palette_np[idx]
                    
                    # 应用最近的颜色
                    output_img[i, j] = best_color
            
            # 确保最终图像只使用调色板中的颜色（使用向量化操作加速）
            output_img_reshaped = output_img.reshape(-1, 3)
            distances = np.sum((output_img_reshaped[:, np.newaxis, :] - palette_np[np.newaxis, :, :]) ** 2, axis=2)
            indices = np.argmin(distances, axis=1)
            output_img = palette_np[indices].reshape(output_img.shape)
        
        return output_img.astype(np.uint8)
    
    def process_image(self, img_path: str, num_colors: int = 2, bayer_n: int = 4, 
                      use_color: bool = False, output_width: int = None, output_height: int = None):
        """
        处理图像并应用Bayer抖动算法。
        
        :param img_path: 图像文件路径
        :param num_colors: 色值数量，默认为2（黑白）
        :param bayer_n: Bayer矩阵阶数
        :param use_color: 是否处理彩色图像，默认为False
        :param output_width: 输出图像宽度，如果为None则保持原始宽度
        :param output_height: 输出图像高度，如果为None则保持原始高度
        :return: 处理后的图像
        """
        # 读取图像
        original_img = cv2.imread(img_path)
        
        if original_img is None:
            raise ValueError(f"无法读取图像: {img_path}")
            
        # 调整图像大小（如果指定了输出尺寸）
        if output_width is not None or output_height is not None:
            original_img = self.resize_image(original_img, output_width, output_height)
        
        # 生成Bayer矩阵
        bayerM = ti.ndarray(shape=(2**bayer_n, 2**bayer_n), dtype=ti.u8)
        bayerM.from_numpy(np.array(self.bayer_matrix(bayer_n), dtype=np.uint8))
        
        if use_color:
            # 彩色图像处理
            self.color_mode = 'color'
            
            # 从图像中提取调色板
            color_palette = self.generate_color_palette_from_image(original_img, num_colors)
            
            if self.DEBUG:
                print(f"已从图像中提取 {num_colors} 种颜色: {color_palette}")
            
            # 制作一个副本用于处理
            img = original_img.copy()
            
            # 应用彩色Bayer抖动
            img = self.apply_color_bayer(img, bayerM, color_palette)
            
            return img
            
        else:
            # 灰度图像处理
            self.color_mode = 'grayscale'
            
            # 生成灰度调色板
            self.generate_palette(num_colors)
            
            # 转换为灰度图
            img = cv2.cvtColor(original_img, cv2.COLOR_BGR2GRAY)
            img = img.astype(np.int32)  # 转换为 int32 类型
            
            # 应用Bayer抖动
            if num_colors > 2:
                # 使用多色版本
                img = self.apply_multi_bayer(img, bayerM, num_colors)
            else:
                # 使用内核版本（黑白）
                self.dither_bayer(img, bayerM)
            
            # 转换回uint8
            img = np.clip(img, 0, self.MAX)
            img = img.astype(np.uint8)
            
            # 确保所有像素值严格使用调色板中的颜色
            palette_np = np.array(self.PALETTE, dtype=np.uint8)
            
            # 使用Taichi内核加速颜色映射
            if self.arch == ti.gpu:
                # 转换为Taichi数组
                img_ti = ti.ndarray(shape=img.shape, dtype=ti.u8)
                img_ti.from_numpy(img)
                
                palette_ti = ti.ndarray(shape=palette_np.shape, dtype=ti.u8)
                palette_ti.from_numpy(palette_np)
                
                # 使用Taichi内核加速映射
                self.map_to_palette(img_ti, palette_ti)
                
                # 转回numpy数组
                img = img_ti.to_numpy()
            else:
                # 使用向量化操作加速CPU处理
                img_flat = img.flatten()
                idx = np.argmin(np.abs(img_flat.reshape(-1, 1) - palette_np.reshape(1, -1)), axis=1)
                img_flat[:] = palette_np[idx]
                img = img_flat.reshape(img.shape)
            
            return img 